fix(memory): roll back every failed SQLite session write - #4203
fix(memory): roll back every failed SQLite session write#4203abhay-codes07 wants to merge 2 commits into
Conversation
openai#4163 established that a write failing partway through leaves an open transaction on the cached connection, and that an open write transaction holds the SQLite write lock for the lifetime of that connection and blocks every later writer. That fix reached only SQLiteSession.add_items. The same defect remained in SQLiteSession.pop_item and clear_session, and in all three AsyncSQLiteSession write paths. clear_session is the clearest case: it issues two DELETEs under one commit, so a failure on the second leaves the first applied inside an open transaction. AsyncSQLiteSession is the most damaging, because it holds one connection for the whole session, so the lock stays held until the session is closed. The rollback obligation belongs to the connection rather than to any single method, so add a _rollback_on_failure(conn) guard per module and apply it at every _locked_connection() write site, including the add_items path that previously inlined it. Commit points are unchanged; only the failure path differs. The guard catches BaseException so an interrupted write cannot strand the lock either.
There was a problem hiding this comment.
Pull request overview
This PR fixes a SQLite transaction-lifecycle bug in the Agents SDK session backends where a failed write could leave an open transaction on a cached/shared connection, stranding the SQLite write lock and blocking subsequent writers. It centralizes rollback-on-failure behavior into a shared guard per module and adds regression tests to ensure failed writes don’t wedge session persistence.
Changes:
- Add a
_rollback_on_failure(...)context guard and apply it to all SQLite session write paths (add_items,pop_item,clear_session) in both sync and async implementations. - Add regression tests that simulate mid-write failures (by dropping tables from a separate connection) and assert the session connection is not left
in_transactionand the write lock is free.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/agents/memory/sqlite_session.py |
Introduces a rollback-on-failure guard and applies it to sync SQLite write operations to avoid stranded write locks. |
src/agents/extensions/memory/async_sqlite_session.py |
Introduces a rollback-on-failure guard and applies it to async SQLite write operations to avoid stranded write locks. |
tests/memory/test_session.py |
Adds regression tests for failed clear_session / pop_item ensuring write lock is released for SQLiteSession. |
tests/extensions/memory/test_async_sqlite_session.py |
Adds regression tests for failed add_items / clear_session / pop_item ensuring write lock is released for AsyncSQLiteSession. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| try: | ||
| yield | ||
| except BaseException: | ||
| await conn.rollback() | ||
| raise |
| try: | ||
| yield | ||
| except BaseException: | ||
| conn.rollback() | ||
| raise |
Review follow-up. Rollback is cleanup, so a connection that is already closed or otherwise unusable would previously replace the failure the caller needs to see. Attempt it best-effort and always re-raise the original. Also add a regression test for cancellation mid-write, which is the case the guard most needs to cover on the async backend: the session holds one connection, so a transaction left open by a cancelled write holds the SQLite write lock until the session is closed.
|
Thanks both review points addressed in 1cf9707, though I only took one of them. Rollback masking the original error (both modules): fixed. Rollback is cleanup, so a connection that is already closed or otherwise unusable would previously replace the failure the caller actually needs. It is now attempted best-effort and the original exception is always re-raised.
The reasoning: by the time the handler runs, I would rather leave it out than add an unexercised branch, but if you know of a path where the rollback await is itself cancelled a second Current state: 6 regression tests fail on |
|
Thanks for sharing this patch. We'll close this PR in favor of #4212, which covers all similar patterns across the SDK. |
Summary
#4163 established that a write failing partway through leaves an open transaction on the cached connection, and that "an open write transaction would hold the SQLite write lock for the lifetime of the connection and block every later writer." That fix reached only
SQLiteSession.add_items.The same defect remained in five sibling write paths. Probing each one — drop a table from an independent connection so a statement inside the write fails, then check the session's connection and whether an independent writer can still take the lock:
clear_sessionis the clearest case: twoDELETEs under one commit, so a failure on the second leaves the first applied inside an open transaction — a partial mutation and a stranded lock.AsyncSQLiteSessionis the most damaging, because it holds one connection for the entire session, so the lock stays held until the session object is closed and one transientOperationalErrorpermanently wedges session persistence for that process.The rollback obligation belongs to the connection rather than to any individual method, so rather than inlining the same
try/except/rollbacksix times this adds one_rollback_on_failure(conn)guard per module and applies it at every_locked_connection()write site — including theadd_itemspath that previously inlined it, which keeps a single source of truth for the concern.Scope notes:
BaseExceptionrather thanException, so an interrupted write cannot strand the lock either. This is the one intentional widening relative to the fix(memory): roll back a failed SQLiteSession insert #4163 inline version.Test plan
Five regression tests, covering each still-broken path:
tests/memory/test_session.py—..._failed_clear_session_releases_write_lock,..._failed_pop_item_releases_write_locktests/extensions/memory/test_async_sqlite_session.py—test_failed_add_items_releases_write_lock,test_failed_clear_session_releases_write_lock,test_failed_pop_item_releases_write_lockEach asserts the connection is no longer
in_transactionand that an independent writer can still take the lock, usingsqlite3.connect(..., timeout=0)to disable the busy handler so a held lock fails immediately rather than stalling. The asyncadd_itemstest reuses the unserializable-item trigger from the #4163 test so the two match, and also asserts the session stays usable afterwards.All five fail on
mainand pass with the fix. The existing #4163 test passes in both runs, which is the control that the probe is measuring the right thing:Verification from the repository root:
make formatmake lintmake mypymain, none in the touched filesmake pyrightmain(src/agents/sandbox/util/tar_utils.py:161)uv run pytest tests/memory/test_session.py tests/extensions/memory/test_async_sqlite_session.pymake testsThe full-suite run was done on Windows, where some sandbox symlink and tracing/realtime timing tests fail independently of this change. I diffed the failing set against a clean
maincheckout in the same environment: the two sets are identical (54 vs 54, no differences either way).Issue number
Closes #4202
Checks
.agents/skills/code-change-verification/scripts/run.sh/reviewbefore submitting this PRThe verification script is a bash script that shells out to
make; I ran the underlying steps individually instead, with the results above.@seratch — this is the follow-through on #4163: I went looking for the same shape in the other write paths and found it in five of them, with
AsyncSQLiteSessionnever having been covered at all.Two judgement calls worth your review. First, I consolidated the inline rollback from #4163 into the shared guard instead of leaving it and adding five more copies — that touches recently merged code, so say the word if you would rather I leave
add_itemsexactly as it is and duplicate the pattern. Second, the guard catchesBaseException;SQLiteSessionruns its writes in a worker thread viaasyncio.to_thread, so cancellation is not the concern there, but an interrupted write stranding the lock is the same failure androllback()is safe in both cases.